import torch
from numpy.typing import ArrayLike


class ProgressRecorder:
    def __init__(self, value: torch.Tensor | ArrayLike = 0):
        try:
            self.process_rank = torch.distributed.get_rank()
            self.process_group_size = torch.distributed.get_world_size()
        except ValueError:
            self.process_rank = 0
            self.process_group_size = 1

        value = torch.asarray(value, dtype=torch.int64)
        self._value = value.detach().sum()

    def get(self) -> int:
        return int(self._value.cpu())

    def increase(self, value: torch.Tensor | ArrayLike = 0):
        value = torch.asarray(value, dtype=torch.int64)
        value = value.detach().sum()
        if self.process_group_size > 1:
            torch.distributed.all_reduce(value, op=torch.distributed.ReduceOp.SUM)
        value = torch.asarray(value, dtype=torch.int64)
        self._value = self._value.to(value.device)
        self._value += value

    def assign(self, value: torch.Tensor | ArrayLike = 0):
        value = torch.asarray(value, dtype=torch.int64)
        self._value = value.detach().sum()
